"""
postgres数据库连接池
"""

import psycopg2
from psycopg2 import pool
from dotenv import load_dotenv
import os

class DBPool:
    def __init__(self,host,port,user,password,dbname):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.dbname = dbname

        self.pool = pool.SimpleConnectionPool(
            minconn=1,
            maxconn=10,
            host=self.host, 
            port=self.port, 
            user=self.user, 
            password=self.password, 
            dbname=self.dbname
        )

    def get_connection(self):
        """
        获取连接
        """
        return self.pool.getconn()

    def put_connection(self, conn):
        """
        归还连接
        """
        self.pool.putconn(conn)

    def close(self):
        """
        关闭连接池
        """
        self.pool.closeall()

load_dotenv(".env")

db_pool = DBPool(
    host=os.getenv("DB_HOST","localhost"),
    port=os.getenv("DB_PORT",5432),
    user=os.getenv("DB_USER","postgres"),
    password=os.getenv("DB_PASSWORD","Zkyy2021"),
    dbname=os.getenv("DB_NAME","duzicong")
)